package simpledb.execution;

import simpledb.common.Type;
import simpledb.storage.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 能够对一组 StringField 字段进行聚合计算的类。
 *
 * 目前仅支持 COUNT 操作；
 * 支持按字段分组（group by）或不进行分组；
 */
public class StringAggregator implements Aggregator {

    private static final long   serialVersionUID = 1L;

    // 存储每个分组的计数结果（Group By 字段 -> 计数值）
    private Map<Field, Integer> groupMap;

    // 分组字段在 Tuple 中的索引，若为 NO_GROUPING 表示不分组
    private int                 gbField;

    // 分组字段的数据类型
    private Type                gbFieldType;

    // 聚合字段在 Tuple 中的索引
    private int                 agField;

    // 使用的聚合操作类型（目前只支持 COUNT）
    private Op                  op;

    // 默认字段（用于无分组时作为键）
    private Field               DEFAULT_FIELD    = new StringField("Default", 10);

    // 缓存输出元组结构描述
    private TupleDesc           td;

    /**
     * 构造函数。
     *
     * @param gbfield 分组字段在 Tuple 中的索引（从 0 开始），若为 NO_GROUPING 表示不分组
     * @param gbfieldtype 分组字段的类型（如 Type.INT_TYPE），若不分组则为 null
     * @param afield 聚合字段在 Tuple 中的索引（从 0 开始）
     * @param what 使用的聚合操作（目前只能是 COUNT）
     * @throws IllegalArgumentException 如果不是 COUNT 操作，则抛出异常
     */
    public StringAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        if (what != Op.COUNT) {
            throw new IllegalArgumentException("字符串聚合器仅支持 COUNT 操作");
        }
        this.groupMap = new HashMap<>();
        this.gbField = gbfield;
        this.agField = afield;
        this.op = what;
        this.gbFieldType = gbfieldtype;
    }

    /**
     * 将一个元组合并进聚合器中。
     * 根据构造函数指定的分组方式，将该元组加入到相应的聚合组中。
     *
     * @param tup 包含聚合字段和分组字段的元组
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        if (this.td == null) {
            buildTupleDesc(tup.getTupleDesc());
        }

        // 获取当前处理的分组字段值
        final Field gbFieldVal = tup.getField(this.gbField);
        // 如果不分组，则使用默认字段作为键
        final Field target = (this.gbField == NO_GROUPING ? DEFAULT_FIELD : gbFieldVal);

        // 增加对应分组的计数值
        this.groupMap.put(target, this.groupMap.getOrDefault(target, 0) + 1);
    }

    /**
     * 构建输出元组的结构描述。
     *
     * 如果没有分组字段，则输出只有一个字段表示 COUNT 结果；
     * 如果有分组字段，则第一个字段是分组字段值，第二个是 COUNT 结果；
     *
     * @param originTd 输入元组的结构描述
     */
    public void buildTupleDesc(final TupleDesc originTd) {
        if (this.gbField == NO_GROUPING) {
            Type[] types = new Type[] { Type.INT_TYPE };
            String[] names = new String[] { "COUNT(*)" };
            this.td = new TupleDesc(types, names);
        } else {
            Type[] types = new Type[] { this.gbFieldType, Type.INT_TYPE };
            String[] names = new String[] { originTd.getFieldName(this.gbField),
                    "COUNT(" + originTd.getFieldName(this.agField) + ")" };
            this.td = new TupleDesc(types, names);
        }
    }

    /**
     * 创建一个迭代器，遍历所有分组的聚合结果。
     *
     * @return 迭代器对象：
     *         - 如果有分组，则返回 (groupVal, countVal) 的元组；
     *         - 如果无分组，则返回单个字段 (countVal)；
     */
    public OpIterator iterator() {
        final List<Tuple> tuples = new ArrayList<>();

        if (this.gbField != NO_GROUPING) {
            this.groupMap.forEach((key, cnt) -> {
                final Tuple tuple = new Tuple(this.td);
                tuple.setField(0, key);
                tuple.setField(1, new IntField(cnt));
                tuples.add(tuple);
            });
        } else {
            final Tuple tuple = new Tuple(this.td);
            tuple.setField(0, new IntField(this.groupMap.get(DEFAULT_FIELD)));
            tuples.add(tuple);
        }

        return new TupleIterator(this.td, tuples);
    }
}
